import glob

import torch
from torch.autograd import Variable
import torchvision.transforms as transforms
from torch.utils.data import Dataset
import numpy as np
import cv2

from net import Net
from netAttention import NetAttention
# net = Net()
net = NetAttention()
right = 0
all = 0
net.eval()
net.cuda()
for i in glob.glob('./test/*.jpg'):
    all+=1
    img = cv2.imread(i)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    img = cv2.resize(img, (256, 256))
    # np.reshape(img,(1,3,256,256))
    net.load_state_dict(torch.load('./model/net_attention28.pth'))
    # net.load_state_dict(torch.load('./model/net28.pth'))
    T = transforms.ToTensor()
    img = T(img)
    img = img.unsqueeze(0)
    img = img.cuda()
    result=net(img)
    result = int(torch.argmax(result))
    truth = int(i.split('\\')[1].split('_')[0])
    print(i,truth,result)
    if truth == result:
        right += 1
print("right:{}, all:{}, acc:{}".format(right,all,right/all))
